package com.k.kc.util.tree.demo4;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.k.kc.util.json.jackson.JacksonUtils;
import com.k.kc.util.tree.demo3.test.OrgEntity;

import java.util.*;
import java.util.stream.Collectors;

public class TreeUtil<T extends TreeUtil.BaseTreeNode> {
    private final List<T> treeList;
    Map<String, TreeUtil.BaseTreeNode> treeMap;

    public TreeUtil(List<T> list) {
        this.treeList = list;
    }

    public static boolean isBlank(CharSequence cs) {
        int strLen = cs == null ? 0 : cs.length();
        if (strLen != 0) {
            for (int i = 0; i < strLen; ++i) {
                if (!Character.isWhitespace(cs.charAt(i))) {
                    return false;
                }
            }
        }
        return true;
    }

    private boolean existParent(TreeUtil.BaseTreeNode node) {
        if (isBlank(node.getParentKey()) || node.getParentKey().equals(node.getKey()) || "-1".equals(node.getParentKey())) {
            return false;
        }
        if (treeMap != null) {
            return treeMap.get(node.getParentKey()) != null;
        }
        return treeList.stream().anyMatch(a -> a.getKey().equals(node.getParentKey()));
    }

    private <E extends TreeUtil.BaseTreeNode> boolean existNodeByKey(E node, String key) {
        return !isBlank(key) && node.getKey().equals(key);
    }

    private <E extends TreeUtil.BaseTreeNode> boolean existNodeByParentKey(E node, String parentKey) {
        return !isBlank(parentKey) && node.getParentKey().equals(parentKey);
    }

    private <E extends TreeUtil.BaseTreeNode> boolean isRootNode(E node) {
        return !existParent(node);
    }

    private <E extends TreeUtil.BaseTreeNode> List<E> makeTree(List<E> allData, String rootKey, String rootParentKey) {
        // 区分了根节点 和 非根节点的数据
        Map<Boolean, List<E>> result = allData.stream()
                .sorted(Comparator.comparingInt(TreeUtil.BaseTreeNode::getSort))
                .collect(Collectors.partitioningBy(
                        a -> (isBlank(rootKey) && isBlank(rootParentKey)) ? isRootNode(a)
                                : existNodeByKey(a, rootKey) || existNodeByParentKey(a, rootParentKey), // 这是你的筛选条件
                        Collectors.toList() // 收集结果的方式
                ));
        List<E> roots = result.get(true);
        List<E> children = result.get(false);
        for (E root : roots) {
            makeChildren(root, children);
        }
        return roots;
    }

    private <E extends TreeUtil.BaseTreeNode> void makeChildren(E root, List<E> nodeList) {
        for (E e : nodeList) {
            if (e.getParentKey().equals(root.getKey())) {
                root.getChildren().add(e);
                makeChildren(e, nodeList);
            }
        }
    }

    /**
     * 指定根节点
     *
     * @param rootKey       rootKey 有多个根节点可以不指定
     * @param rootParentKey rootParentKey
     * @return result
     */
    public List<T> build(String rootKey, String rootParentKey) {
        if (treeList == null || treeList.isEmpty()) {
            throw new IllegalArgumentException("argument error");
        }
        // key 重复会抛出异常
        treeMap = treeList.parallelStream().collect(Collectors.toMap(TreeUtil.BaseTreeNode::getKey, n -> n));
        if (treeMap.size() > 0) {
            return makeTree(treeList, rootKey, rootParentKey);
        }
        return Collections.emptyList();
    }

    /**
     * 不指定根节点
     *
     * @return Result
     */
    public List<T> build() {
        return build(null, null);
    }
    public static class BaseTreeNode {
        private String key;
        private String parentKey;
        private List<TreeUtil.BaseTreeNode> children = new ArrayList<>();
        private boolean isLeaf;
        private int sort;

        public String getKey() {
            return key;
        }

        public void setKey(String key) {
            this.key = key;
        }

        public String getParentKey() {
            return parentKey;
        }

        public void setParentKey(String parentKey) {
            this.parentKey = parentKey;
        }

        public List<TreeUtil.BaseTreeNode> getChildren() {
            return children;
        }

        public void setChildren(List<TreeUtil.BaseTreeNode> children) {
            this.children = children;
        }

        public boolean getLeaf() {
            return isLeaf;
        }

        public void setLeaf(boolean leaf) {
            isLeaf = leaf;
        }

        public int getSort() {
            return sort;
        }

        public void setSort(int sort) {
            this.sort = sort;
        }
    }

    public static void main(String[] args) throws JsonProcessingException {
        TreeUtil<OrgTreeDTO> treeUtil = new TreeUtil<>(convertData(getTestData()));
        List<OrgTreeDTO> treeList = treeUtil.build("101", "-1");
        System.out.println(JacksonUtils.toJson(treeList));

    }

    public static List<OrgEntity> getTestData() {
        List<OrgEntity> list = new ArrayList<>();
        // root
        OrgEntity root1 = new OrgEntity("1", "101", "黑龙江", "-1", 3);
        OrgEntity root2 = new OrgEntity("2", "102", "吉林", "-1", 2);
        OrgEntity root3 = new OrgEntity("3", "103", "辽宁", "-1", 1);

        // 市
        OrgEntity rootChild11 = new OrgEntity("11", "101-01", "哈尔滨", "101", 6);
        OrgEntity rootChild12 = new OrgEntity("12", "101-02", "齐齐哈尔", "101", 5);

        // 区
        OrgEntity rootChild1101 = new OrgEntity("111", "101-01-01", "松北区", "101-01", 7);
        list.add(root1);
        list.add(root2);
        list.add(root3);
        list.add(rootChild11);
        list.add(rootChild12);
        list.add(rootChild1101);
        return list;
    }

    public static List<OrgTreeDTO> convertData(List<OrgEntity> list) {
        return list.stream().map(a -> {
            OrgTreeDTO treeDTO = new OrgTreeDTO();
            // copy original
            treeDTO.setId(a.getId());
            treeDTO.setCode(a.getCode());
            treeDTO.setName(a.getName());
            treeDTO.setParentCode(a.getParentCode());

            // set Base
            treeDTO.setKey(a.getCode());
            treeDTO.setParentKey(a.getParentCode());
            treeDTO.setSort(a.getSort() == null ? 1 : a.getSort());
            return treeDTO;
        }).collect(Collectors.toList());
    }
}
